"""当选出某种不变量时，验证其是否是符合要求的"""

from typing import Iterable
import argparse
import os
import numpy as np
import yaml
import tqdm
from itertools import combinations


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )
    arg_parser.add_argument(
        "-o",
        "--catalog_dir",
        type=str,
        default="/disk527/sdb1/a804_cbf/catalog",
        help="The directory of the catalog",
    )
    arg_parser.add_argument(
        "--center",
        type=float,
        default=0.5,
    )
    arg_parser.add_argument(
        "--edge",
        type=float,
        default=0.5,
    )
    arg_parser.add_argument(
        "--sigma",
        type=float,
        default=0.3,
    )
    return arg_parser.parse_args()


def noiser(points, center: float = 0.5, edge: float = 0.5, sigma: float = 1):
    """处理中心偏离和边缘偏离"""
    if np.random.rand() < center:
        points = points + np.random.normal(0, sigma)
    if np.random.rand() < edge:
        points = points + np.random.normal(0, sigma, size=points.shape)
    return points


def load_for_match(points_id_path: str, *args) -> Iterable:
    """
    此处加入噪声，主要分为四种：中心偏离、边缘偏离、误增加、误减少
    """
    with open(points_id_path, "r") as f:
        f.readline()
        points_2d = []
        points_id = []
        for line in map(lambda x: x.strip().split(","), f):
            idx = line[-1]
            if len(points_id) == 0:
                points_id.append(idx)
                points_2d.append(list(map(float, line[:2])))
            elif idx != points_id[-1]:
                yield points_id[-1], noiser(np.array(points_2d), *args)
                points_id.append(idx)
                points_2d = []
            else:
                points_2d.append(list(map(float, line[:2])))
        yield points_id[-1], noiser(np.array(points_2d), *args)


def preload(catalog_dir):
    """
    Read the catalog file and store it in the memory
    Arguments:
        catalog_dir (str) : the path to the catalog file directory. The catalog file should have a similar format to the following:
            ```
            name,x,y,z,param1,param2,param3,param4,param5
            D1,12,31,141,15,1,1,1,1
            D2,2,2,2,2,2,2.3,2.2,2
            ...
            ```
    """
    catalog = []
    catalog_name = []
    with open(os.path.join(catalog_dir, "catalog.csv"), "r") as f:
        f.readline()  # 去除标题
        for line in f:
            catalog.append(list(map(float, line.split(",")[2:])))
            catalog_name.append(line.split(",")[1])
    catalog = np.array(catalog)
    catalog_name = np.array(catalog_name)
    # 加载catalog的其他信息
    with open(os.path.join(catalog_dir, "config.yaml"), "r") as f:
        catalog_info = yaml.load(f, Loader=yaml.FullLoader)
    return catalog, catalog_name, catalog_info


def fit_ellipse(coor, return_error=False):
    """
    拟合椭圆
    Returens:
        C (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y)
    """
    x, y = coor
    A = np.vstack((x**2, x * y, y**2, x, y, np.ones_like(x))).T
    *_, Vh = np.linalg.svd(A)
    C = Vh[-1]
    if return_error:
        error = np.sqrt(np.mean((A @ C) ** 2))
        return C, error
    return C


def ellipse_cnt(ellipse):
    """
    从椭圆参数中还原椭圆的参数
    """
    # calculate the centroid using the ellipse
    cnt_x = (ellipse[1] * ellipse[4] - 2 * ellipse[2] * ellipse[3]) / (
        4 * ellipse[0] * ellipse[2] - ellipse[1] ** 2
    )
    cnt_y = (ellipse[1] * ellipse[3] - 2 * ellipse[0] * ellipse[4]) / (
        4 * ellipse[0] * ellipse[2] - ellipse[1] ** 2
    )

    return cnt_x, cnt_y


def invariant(param1, param2):
    """
    需要测试的不变量计算函数
    Arguments:
        param1 (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1), 6 x N array
        param2 (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1), 6 x N array
    """
    C1 = np.array(
        [
            [param1[0], param1[1] / 2, param1[3] / 2],
            [param1[1] / 2, param1[2], param1[4] / 2],
            [param1[3] / 2, param1[4] / 2, param1[5]],
        ]
    )
    C2 = np.array(
        [
            [param2[0], param2[1] / 2, param2[3] / 2],
            [param2[1] / 2, param2[2], param2[4] / 2],
            [param2[3] / 2, param2[4] / 2, param2[5]],
        ]
    )
    cnt1, cnt2 = np.array(ellipse_cnt(param1)), np.array(ellipse_cnt(param2))
    cnt1 = np.pad(cnt1, ((0, 1)), "constant", constant_values=1)
    cnt2 = np.pad(cnt2, ((0, 1)), "constant", constant_values=1)
    # 直接除以行列式是不对滴，参见文献
    # [1] Gros P. , Quan L. .Projective invariants for Vision[R].1992:47
    # det1 = np.linalg.det(C1)
    # det2 = np.linalg.det(C2)

    # C1 = C1 / np.power(np.abs(det1), 1 / 3) * np.sign(det1)
    # C2 = C2 / np.power(np.abs(det2), 1 / 3) * np.sign(det2)
    I1 = (cnt1 @ C2 @ cnt1) / (cnt2 @ C2 @ cnt2)
    I2 = (cnt2 @ C1 @ cnt2) / (cnt1 @ C1 @ cnt1)
    # I1 = np.trace(C1 @ np.linalg.inv(C2))
    # I2 = np.trace(C2 @ np.linalg.inv(C1))
    # 用指数函数放大不变量
    # I1 = np.cosh(I1)
    # I2 = np.cosh(I2)
    return I1, I2


if __name__ == "__main__":
    args = arg_parse()
    points_dir = os.path.join(args.collect_dir, "points_id")
    catalog, catalog_name, catalog_info = preload(args.catalog_dir)
    output_file = "_".join(
        f"{k}_{v}" for k, v in args.__dict__.items() if "dir" not in k
    )
    error_by_noise = []
    it = 0
    with open(f"{args.collect_dir}/pose.csv", "r") as f:
        f.readline()  # 去除标题
        file_names = sorted(
            os.listdir(points_dir), key=lambda x: float(x.split(".")[0])
        )
        for time in tqdm.tqdm(
            map(lambda x: x.split(".txt")[0], file_names), total=len(file_names)
        ):
            if float(time) < 100:
                continue
            for line in f:
                if abs(float(line.split(",")[0]) - float(time)) < 2e-4:
                    pose = [float(x) for x in line.strip().split(",")[1:]]
                    break
                elif float(line.split(",")[0]) > float(time):
                    pose = [float(x) for x in line.strip().split(",")[1:]]
                    break
            ellipsis = []
            gt_id = []
            try:
                for idx, points in load_for_match(
                    os.path.join(points_dir, time + ".txt"),
                    args.center,
                    args.edge,
                    args.sigma,
                ):
                    elli, err = fit_ellipse(np.array(points).T, return_error=True)
                    if err < 1e-3:
                        ellipsis.append(elli)
                        gt_id.append(idx)
            except:
                continue
            it += 1
            for i, j in combinations(range(len(gt_id)), 2):
                # 导航坑表内的不变量
                I, J = gt_id[i], gt_id[j]
                I_ind = np.where(catalog_name == I)[0][0]
                J_ind = np.where(catalog_name == J)[0][0]
                I1, I2 = invariant(catalog[I_ind][3:], catalog[J_ind][3:])

                # 图像导航坑计算出的不变量
                i1, i2 = invariant(ellipsis[i], ellipsis[j])
                # print(f"real invariant: {i1}, {i2}, calculated invariant: {I1}, {I2}")
                if np.isnan((I1 - i1) / I1) or np.isnan((I2 - i2) / I2):
                    continue
                if I == "79013" or J == "79013":
                    continue
                error_by_noise.append(((I1 - i1) / I1, (I2 - i2) / I2))
            if it % 100:
                output = np.array(error_by_noise)
                ind = (np.abs(output[:, 0]) < 1) & (np.abs(output[:, 1]) < 1)
                print(f"mean error: {np.sqrt(np.mean(np.sum(output[ind]**2, axis=1)))}")
# for C-1C2 error is 0.154328
# for (x1 C x2) ** 2 / (x1 C x1) / (x2 C x2) error is 0.31928
# for det C1 = det C2 = 1, (x1 C1 x2) / (x2 C2 x1) error is 0.16145
# for (x2 C2 x2) /(x1 C1 x2) (x2 C1 x2) error is 0.1584
